#* 这个是用于 编译pta优化的 （arm）

# docker run -it -v /Users/x-contion/AscendCode/Huawei:/data --name pta_arm_compiler quay.io/pypa/manylinux2014_aarch64:latest
# docker run -it -v /mnt/d/AscendCode/Huawei:/data --platform linux/aarch64 --name pta_arm_compiler quay.io/pypa/manylinux2014_aarch64:latest

# docker run -it -v /data/torch_compile:/data --platform linux/aarch64 --name pta_arm_compiler quay.io/pypa/manylinux2014_aarch64:latest

flag_use_pgo=1 #* 0表示不使用，1 表示生成 pgo 优化，2 表示使用 pgo 优化
## 2 编译python
export CC=clang
export C=clang
export CXX=clang++

if [[ ! -d /tmp/Python ]];then
    echo "复制python源码到 /tmp/Python"
    mkdir -p /tmp/Python
    cp $path_resource/python/Python-3.8.20.tgz /tmp/Python/
    tar -xvf $path_resource/python/Python-3.8.20.tgz -C /tmp/Python
    #* 还要手动编译一下 libffi
    cp -r $path_resource/libffi /tmp/libffi
    cd /tmp/libffi
    ./configure --prefix=/usr/local/libffi
    make -j 32
    make install
    echo "export LD_LIBRARY_PATH=/usr/local/libffi/lib:\$LD_LIBRARY_PATH" >> ~/.bashrc

    #* 补充一些工具包
    cd /tmp/Python/Python-3.8.20
    ./configure --prefix=/usr/local/python3.8 --with-system-ffi --with-lto --enable-optimizations
    make -j 32
    make install
    echo "替换本地的python"
    echo "export PATH=/usr/local/python3.8/bin:\$PATH" >> ~/.bashrc

    # rm -f /usr/bin/python
    rm -f /usr/bin/python3
    rm -f /usr/bin/pip
    rm -f /usr/bin/pip3

    ln -s /usr/local/python3.8/bin/python3 /usr/bin/python3
    # ln -s /usr/local/python3.8/bin/python3 /usr/bin/python
    ln -s /usr/local/python3.8/bin/pip3 /usr/bin/pip
    ln -s /usr/local/python3.8/bin/pip3 /usr/bin/pip3
else
    echo "已经存在python源码, 不再重复编译"
fi

## 3 编译pytorch
#? pip config set global.index-url http://mirrors.ayers.iflytek.cn/pypi/web/simple/
#? pip config set global.extra-index-url http://mirrors.ayers.iflytek.cn:9898/
#? pip config set install.trusted-host mirrors.ayers.iflytek.cn

cd $path_resource
if [[ ! -d /tmp/pytorch ]];then
    echo "复制pytorch源码到 /tmp/pytorch"
    cp -r pytorch/pytorch_v2.1 /tmp/pytorch
fi
cd /tmp/pytorch
#! 切换到 flto 或 pgo 优化
cat /tmp/pytorch/CMakeLists.txt | grep "string(APPEND\ CMAKE_CXX_FLAGS\ \"\ -O2" # 检查

if [[ $flag_use_pgo == 0 ]];then
    string_new_cmd="-mcpu=tsv110\ -fPIC\ -flto=thin\ -fuse-ld=lld"
fi
if [[ $flag_use_pgo == 1 ]];then
    string_new_cmd="-mcpu=tsv110\ -fPIC\ -flto=thin\ -fuse-ld=lld -fprofile-generate=/root/pgo_profile"
fi
if [[ $flag_use_pgo == 2 ]];then
    string_new_cmd="-mcpu=tsv110\ -fPIC\ -flto=thin\ -fuse-ld=lld -fprofile-use=/root/pgo_profile/pgo.profdata"
fi

sed -i "s@string(APPEND\ CMAKE_CXX_FLAGS\ \"\ -O2\ .*@string(APPEND\ CMAKE_CXX_FLAGS\ \"\ -O2\ $string_new_cmd\")@g" /tmp/pytorch/CMakeLists.txt
cat /tmp/pytorch/CMakeLists.txt | grep "string(APPEND\ CMAKE_CXX_FLAGS\ \"\ -O2" # 检查

string_new_cmd="append_cxx_flag_if_supported(\"-Werror=cast-function-type\" CMAKE_CXX_FLAGS)"
sed -i "s@^.*${string_new_cmd}.*\$@\ \ #\ ${string_new_cmd}@g" /tmp/pytorch/CMakeLists.txt
cat /tmp/pytorch/CMakeLists.txt | grep 'Werror=cast-function-type'
/usr/local/bin/python3.8 -m pip install pyyaml
/usr/local/bin/python3.8 setup.py bdist_wheel
echo "编译后的位置在 dist/torch*.whl"
# 编译 torch 后安装依赖
cd /tmp/pytorch
/usr/local/bin/python3.8 -m pip install -r requirements.txt
export USE_XNNPACK=0
/usr/local/bin/python3.8 -m pip install dist/torch*.whl

if [[ $flag_use_pgo == 0 ]];then
    target_dir=/data/torch_compile/flto
fi
if [[ $flag_use_pgo == 1 ]];then
    target_dir=/data/torch_compile/flto_and_pgo_1st
fi
if [[ $flag_use_pgo == 2 ]];then
    target_dir=/data/torch_compile/flto_and_pgo_2nd
fi
mkdir -p $target_dir
rm -f $target_dir/torch*.whl
cp -f dist/torch*.whl $target_dir/

## 3 编译pta
cd $path_resource
if [[ ! -d /tmp/pta_v2.1.0 ]];then
    echo "复制pta源码到 /tmp/pta_v2.1.0"
    cp -r pta/pta_v2.1.0 /tmp/pta_v2.1.0
fi
cd /tmp/pta_v2.1.0

cat /tmp/pta_v2.1.0/CMakeLists.txt | grep -n "SET(CMAKE_CXX_FLAGS_RELEASE\ \"-O2"
if [[ $flag_use_pgo == 0 ]];then
    string_new_cmd="-mcpu=tsv110\ -fGNU-compatibility\ -flto=thin\ -fuse-ld=lld\ "
fi
if [[ $flag_use_pgo == 1 ]];then
    string_new_cmd="-mcpu=tsv110\ -fGNU-compatibility\ -flto=thin\ -fuse-ld=lld\ -fprofile-generate=/root/pgo_profile"
fi
if [[ $flag_use_pgo == 2 ]];then
    string_new_cmd="-mcpu=tsv110\ -fGNU-compatibility\ -flto=thin\ -fuse-ld=lld\ -fprofile-use=/root/pgo_profile/pgo.profdata"
fi
sed -i "s@^SET(CMAKE_CXX_FLAGS_RELEASE\ \"-O2.*@SET(CMAKE_CXX_FLAGS_RELEASE\ \"-O2\ ${string_new_cmd}\")@g" /tmp/pta_v2.1.0/CMakeLists.txt
cat /tmp/pta_v2.1.0/CMakeLists.txt | grep -n "SET(CMAKE_CXX_FLAGS_RELEASE\ \"-O2"  # 检查

cd /tmp/pta_v2.1.0
bash ci/build.sh --python=3.8
echo "编译后的位置在 dist/torch_npu*.whl"

if [[ $flag_use_pgo == 0 ]];then
    target_dir=/data/torch_compile/flto
fi
if [[ $flag_use_pgo == 1 ]];then
    target_dir=/data/torch_compile/flto_and_pgo_1st
fi
if [[ $flag_use_pgo == 2 ]];then
    target_dir=/data/torch_compile/flto_and_pgo_2nd
fi
cp -f dist/torch_npu*.whl $target_dir/